/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.mapreduce.example;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.StringTokenizer;

import org.apache.commons.cli.BasicParser;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.sss.client.DataGetter;
import org.sss.client.DataPutter;
import org.sss.client.SimpleJob;
import org.sss.client.SssClient;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.Mapper;
import org.sss.mapreduce.Output;
import org.sss.mapreduce.Reducer;
import org.sss.mapreduce.SideData;
import org.sss.mapreduce.Tuple;
import org.sss.mapreduce.datatype.PackableBytes;
import org.sss.mapreduce.datatype.PackableInt;
import org.sss.mapreduce.datatype.PackableString;

public class SideDataExample {
  static Class<? extends Mapper>   mapperClass      = SideDataExampleMapper.class;
  static Class<? extends Reducer>  combinerClass    = IntSumReducer.class;
  static Class<? extends Reducer>  reducerClass     = IntSumReducer.class;

  public static GroupID run(SssClient client, GroupID input, GroupID sidedata) throws Exception {
    SimpleJob job = new SimpleJob(client);

    job.setMapperClass(mapperClass);
    job.setCombinerClass(combinerClass);
    job.setReducerClass(reducerClass);
    job.setInputGID(input);
    job.getConfiguration().set("sidedataexample.sidedata", sidedata.toString());

    return job.execute();
  }

  /**
   * Usage:
   */
  public static void main(String[] args) throws Exception {
    SssClient client = new SssClient(args);
    long start;
    long end;
    long elapsed;
    int n = 0;
    boolean broadcast = false;

    String[] others = client.getArgs();
    BasicParser parser = new BasicParser();
    Options options = new Options();
    options.addOption("b", false, "broadcast sidedata.");
    try {
      CommandLine line = parser.parse(options, others);

      if (line.getArgs().length < 1) {
        System.err.println("no recod size");
        System.exit(1);
      }
      broadcast = line.hasOption("b");
      n = Integer.parseInt(line.getArgs()[0]);
    }
    catch(ParseException e) {
      System.err.println(e.getMessage());
      HelpFormatter f = new HelpFormatter();
      f.printHelp("SideDataExample", options);
      System.exit(1);
    }

    PackableBytes val;
    {
      StringBuilder sb = new StringBuilder();
      for (char c = 'a';c <= 'z';++c) {
        for (int i = 0;i < 5;++i) {
          for (int j = 0;j < 5;++j) {
            sb.append(c);
          }
          sb.append(' ');
        }
        sb.append('\n');
      }
      val = new PackableBytes(sb.toString().getBytes("UTF-8"));
    }

    // Upload Side data
    start = System.currentTimeMillis();
    GroupID sidedata = GroupID.createRandom(client, broadcast, false);
    {
      DataPutter<PackableString, PackableString> putter =
        DataPutter.create(client, PackableString.class, PackableString.class, sidedata);
      try {
        for (char c = 'a';c <= 'z';++c) {
          StringBuilder sb = new StringBuilder();
          for (int i = 0;i < 5;++i) {
            sb.append(c);
          }
          String s = sb.toString();
          putter.append(new PackableString(s),
                        new PackableString(s.toUpperCase()));
        }
      } finally {
        putter.close();
      }
    }

    // Upload
    GroupID table = null;
    {
      DataPutter<PackableBytes, PackableBytes>  putter =
        DataPutter.create(client, PackableBytes.class, PackableBytes.class);
      table = putter.getOutputGroupID();
      try {
        for (int i = 0;i < n;) {
          int stop  = Math.min(i + 4096, n);
          for (;i < stop;++i) {
            putter.append(
                  new PackableBytes(Integer.toString(i).getBytes("UTF-8")), val);
          }
        }
      } finally {
        putter.close();
      }
    }

    end = System.currentTimeMillis();
    elapsed = end - start;

    System.out.println("Data upload(" + table + ") elapsed " + elapsed + " ms.");
    System.out.println("SideData GroupID = " + sidedata);

    // Word Count
    start = System.currentTimeMillis();
    GroupID out = run(client, table, sidedata);
    end = System.currentTimeMillis();
    elapsed = end - start;

    System.out.println("Word count elapsed " + elapsed + " ms." + out.toString());

    // Get Result
    start = System.currentTimeMillis();

    Map<String, Integer> result = new HashMap<String, Integer>();
    DataGetter<PackableString, PackableInt> getter =
      DataGetter.create(client, PackableString.class, PackableInt.class, out);
    try {
      for (Tuple<PackableString, PackableInt> kv: getter) {
        String word = kv.getKey().get();
        int cnt = kv.getValue().get();

        if (result.containsKey(word)) {
          System.err.println(word + " duplicate.");
        } else {
          result.put(word, cnt);
        }
      }
    } finally {
      getter.close();
    }

    end = System.currentTimeMillis();
    elapsed = end - start;

    System.out.println("\"Gets result\" elapsed " + elapsed + " ms.");

    // Check Result
    Map<String, Integer> expect = new HashMap<String, Integer>();
    for (char c = 'A';c <= 'Z';++c) {
      StringBuilder sb = new StringBuilder();
      for (int i = 0;i < 5;++i) {
        sb.append(c);
      }
      expect.put(sb.toString(), n * 5);
    }

    if (result.equals(expect)) {
      System.out.println("Expect result.");
    } else {
      System.err.println("Unexpect result.");
    }
  }

  public static class SideDataExampleMapper extends Mapper {
    private SideData<PackableString, PackableString> sidedata = null;

    @Override
    public void configure(Context context) throws Exception {
      Configuration config = context.getConfiguration();

      GroupID sgid = GroupID.createFromString(config.get("sidedataexample.sidedata"));

      sidedata = context.openSideData(sgid, PackableString.class, PackableString.class);
    }

    private void tokenizerMap(byte[] value, Output<PackableString, PackableInt> output) throws Exception {
      StringTokenizer it = new StringTokenizer(new String(value));
      while (it.hasMoreTokens()) {
        PackableString str = new PackableString(it.nextToken());
        List<PackableString> sd = sidedata.get(str);
        if (sd != null && sd.size() == 1) {
          str = sd.get(0);
        }
        output.write(str, new PackableInt(1));
      }
    }
    public void map(Context context,
                    PackableBytes key, PackableBytes value,
                    Output<PackableString, PackableInt> output)
      throws Exception {
      tokenizerMap(value.get(), output);
    }
  }
}
